import numpy as np
import matplotlib.pyplot as plt


class Interpolation:
    def __init__(self, polynomial, x, y, y0):
        self.polynomial = polynomial
        self.x = x
        self.y0 = y0
        self.y = y
        print('Initiated')

    def cal_interp_x0(self, x0):
        """
        计算所给定的插值点的数值，即插值
        :param x0: 所求插值点的坐标
        :return:
        """
        x0 = np.asarray(x0, dtype=np.float)
        n0 = len(x0)  # 所求插值点的个数
        y_0 = np.zeros(n0)  # 存储插值点x0对应的插值
        t = self.polynomial.free_symbols.pop()  # 返回个集合。自由符号：获取前面定义的符号 自由变量
        for i in range(n0):
            y_0[i] = self.polynomial.evalf(subs={t: x0[i]})
        self.y0 = y_0
        return y_0

    def plt_interpolation(self, x0=None, y0=None):
        """
        可视化插值图像和所求插值点
        :return:
        """
        plt.Figure(figsize=(8, 6))
        plt.plot(self.x, self.y, "ro", label="Interpolation base points")
        xi = np.linspace(min(self.x), max(self.x), 100)  # 模拟100个值
        yi = self.cal_interp_x0(xi)
        plt.plot(xi, yi, "b--", label="Interpolatin polynomial")
        if x0 is not None and y0 is not None:
            plt.plot(x0, y0, "g*", label="Interpolation values")
        plt.legend()
        plt.xlabel("x", fontdict={"fontsize": 12})
        plt.ylabel("y", fontdict={"fontsize": 12})
        plt.title("Lagrange interpolation polynoial and values", fontdict={"fontsize": 12})
        plt.grid()
        plt.show()
